#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : cd
# @Email : 19688513@qq.com
# @Project : horizons-engine-pybroker
# @File : HistoryStockDataSource.py
# @Software: PyCharm
# @time: 2025/6/20 11:09
# @description : 历史股票数据源实现
"""

import urllib3
import logging
import pandas as pd
import pybroker
from pybroker.data import DataSource
from pybroker import Strategy

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class HistoryStockDataSource(DataSource):
    def __init__(self, file_path='data/prices.csv'):
        """
        历史股票数据源

        参数:
            file_path: CSV文件路径
        """
        super().__init__()
        # 注册所有自定义列
        pybroker.register_columns(
            'volume', 'turnover', 'amplitude',
            'change_pct', 'change_amt', 'turnover_rate'
        )
        self.file_path = file_path
        self._data_cache = None  # 数据缓存提高性能

    def _load_data(self):
        """加载并预处理历史数据"""
        if self._data_cache is None:
            # 读取CSV文件 - 确保symbol列作为字符串读取
            df = pd.read_csv(
                self.file_path,
                dtype={'symbol': str},  # 确保symbol列是字符串类型
                low_memory=False  # 避免混合类型警告
            )

            # 转换日期格式
            df['date'] = pd.to_datetime(df['date'])

            # 按日期排序确保时间序列正确
            df = df.sort_values('date')

            # 添加唯一标识符 - 确保所有列都是字符串
            df['id'] = df['date'].astype(str) + '_' + df['symbol'].astype(str)

            self._data_cache = df

        return self._data_cache

    def _fetch_data(self, symbols, start_date, end_date, _timeframe, _adjust):
        """
        获取指定范围的历史数据

        参数:
            symbols: 股票代码列表
            start_date: 开始日期
            end_date: 结束日期
            _timeframe: 时间框架（此处未使用）
            _adjust: 复权设置（此处未使用）
        """
        df = self._load_data()

        # 筛选股票代码
        if symbols:
            # 确保symbols是字符串列表
            symbols = [str(sym) for sym in symbols]
            df = df[df['symbol'].isin(symbols)]

        # 筛选日期范围
        start_dt = pd.to_datetime(start_date)
        end_dt = pd.to_datetime(end_date)
        date_mask = (df['date'] >= start_dt) & (df['date'] <= end_dt)

        return df[date_mask]

    def get_full_data(self):
        """获取完整历史数据集（用于分析）"""
        return self._load_data()


# 初始化历史数据源
history_data = HistoryStockDataSource()

# 1. 直接查询历史数据
try:
    df = history_data.query(['000001'], '2004-01-01', '2025-12-01')
    print(f"查询到 {len(df)} 条历史记录")
    if not df.empty:
        print("示例数据:")
        print(df[['date', 'symbol', 'open', 'high', 'low', 'close']].head(3))
except Exception as e:
    print(f"查询历史数据时出错: {str(e)}")


# 2. 在策略中使用历史数据
def volume_based_strategy(ctx):
    # 访问历史成交量数据
    if len(ctx.volume) < 5:
        return

    current_vol = ctx.volume[-1]
    avg_vol = ctx.volume[-5:].mean()  # 5日平均成交量

    if not ctx.long_pos():
        # 成交量突破时买入
        if current_vol > avg_vol * 1.5:
            ctx.buy_shares = 100
    else:
        # 成交量萎缩时卖出
        if current_vol < avg_vol * 0.7:
            ctx.sell_shares = ctx.long_pos().shares


# 创建策略
try:
    strategy = Strategy(
        data_source=history_data,
        start_date='2004-01-01',
        end_date='2005-12-31'
    )
    strategy.add_execution(volume_based_strategy, ['000001'])
    print("策略创建成功!")
except Exception as e:
    print(f"创建策略时出错: {str(e)}")

# 3. 获取完整历史数据集进行分析
try:
    full_history = history_data.get_full_data()
    print(f"完整历史数据集包含 {len(full_history)} 条记录")
    if not full_history.empty:
        print("时间范围:", full_history['date'].min(), "至", full_history['date'].max())
        print("包含的股票代码:", full_history['symbol'].unique()[:5])  # 显示前5个股票代码
except Exception as e:
    print(f"获取完整历史数据时出错: {str(e)}")